# 将工程生成的*.csv文件数据可视化，生成折线图

import pandas as pd
import matplotlib.pyplot as plt
import re
import os

def extract_numeric(value):
    match = re.search(r"(\d+\.\d+|\d+)", value)
    # 如果单位是微秒，除以1000
    if 'us' in value:
        return float(match.group(0)) / 1000 if match else None
    return float(match.group(0)) if match else None

def plot_alloc_free_times(csv_path, x_var_name='size'):
    # 读取CSV文件
    df = pd.read_csv(csv_path)
    
    # 筛选hmalloc和malloc方法的数据
    hmalloc_df = df[df['alloc_func'] == 'hmalloc']
    malloc_df = df[df['alloc_func'] == 'malloc']
    
    infos = [None, None, None]  # 示例初始化，根据实际情况调整

    # 使用列表推导式安全地移除x_var_name
    titles = ['times', 'works', 'rounds', 'size']
    titles = [title for title in titles if title != x_var_name]

    for i, title in enumerate(titles):
        if title in hmalloc_df.columns:
            infos[i] = hmalloc_df[title].iloc[0]

    # print(infos)

    # 读取平均分配和释放时间
    hmalloc_df.loc[:, 'avg_alloc_time'] = hmalloc_df['avg_alloc_time'].apply(extract_numeric)
    hmalloc_df.loc[:, 'avg_free_time'] = hmalloc_df['avg_free_time'].apply(extract_numeric)
    malloc_df.loc[:, 'avg_alloc_time'] = malloc_df['avg_alloc_time'].apply(extract_numeric)
    malloc_df.loc[:, 'avg_free_time'] = malloc_df['avg_free_time'].apply(extract_numeric)

    # 按x_var_name分组计算平均分配和释放时间
    hmalloc_avg = hmalloc_df.groupby(x_var_name)[['avg_alloc_time', 'avg_free_time']].mean()
    malloc_avg = malloc_df.groupby(x_var_name)[['avg_alloc_time', 'avg_free_time']].mean()

    # 绘制图表
    plt.figure(figsize=(10, 5))
    
    # 绘制折线图
    plt.plot(hmalloc_avg.index, hmalloc_avg['avg_alloc_time'], label='hmalloc - Avg Alloc Time', marker='o')
    plt.plot(hmalloc_avg.index, hmalloc_avg['avg_free_time'], label='hmalloc - Avg Free Time', marker='o')
    plt.plot(malloc_avg.index, malloc_avg['avg_alloc_time'], label='malloc - Avg Alloc Time', marker='x')
    plt.plot(malloc_avg.index, malloc_avg['avg_free_time'], label='malloc - Avg Free Time', marker='x')

    plt.xscale('log')   # 设置x轴为对数坐标轴（非线性）
    plt.xlabel(x_var_name)

    plt.yscale('log')   # 设置x轴为对数坐标轴（非线性）
    plt.ylabel('Time (ms)')

    # 标题：titles和infos信息，给出自变量和因变量的含义
    plt.title(f'Comparison of Allocation and Free Time for hmalloc vs malloc \n {titles[0]}: {infos[0]}, {titles[1]}: {infos[1]}, {titles[2]}: {infos[2]}')
    plt.legend()
    plt.grid(True)
    # plt.show()

    # 构建保存图片的路径
    base_file_path = os.path.splitext(csv_path)[0]  # 移除文件后缀
    output_file_path = f"{base_file_path}.png"  # 添加新的文件后缀

    # 使用plt.savefig()代替plt.show()
    plt.savefig(output_file_path)
    # print(f"图表已保存至：{output_file_path}")

if __name__ == '__main__':
    plot_alloc_free_times('./benchmark_size.csv', x_var_name='size')
    plot_alloc_free_times('./benchmark_times.csv', x_var_name='times')
    plot_alloc_free_times('./benchmark_works.csv', x_var_name='works')